package com.wgf.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

@RestController
@Slf4j
public class EmbeddinglController
{
    @Resource
    private EmbeddingModel embeddingModel;//用于将文本转换为向量
    @Resource
    private QdrantClient qdrantClient;//创建Qdrant客户端
    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;//查询向量进行相似性搜索

    /**
     * 文本向量化测试，看看形成向量后的文本
     * http://localhost:9012/embedding/embed     端口号9012被占用，暂时用9112
     * @return
     */
    @GetMapping(value = "/embedding/embed")//用于测试嵌入模型是否正常工作，以及查看文本被转换成什么样的向量
    public String embed(){//作用：将给定的文本转换为向量并返回

        String prompt = """
                   咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);//调用embeddingModel.embed()方法将文本转换为向量

        System.out.println(embeddingResponse);

        return embeddingResponse.content().toString();//embeddingResponse.content()从响应对象中提取实际的嵌入内容,toString()将 Embedding 对象转换为字符串表示形式
    }

    /**
     * 新建向量数据库实例和创建索引：test-qdrant
     * 类似mysql create database test-qdrant
     * http://localhost:9012/embedding/createCollection
     */
    @GetMapping(value = "/embedding/createCollection")//在使用Qdrant存储向量之前，需要先创建一个集合来存储它们。
    public void createCollection()//作用：在Qdrant中创建一个新的集合（类似于数据库中的表）
    {
        var vectorParams = Collections.VectorParams.newBuilder()//异步创建集合，指定集合名称和向量参数（包括距离度量和向量大小）
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)//设置向量参数：使用余弦相似度作为距离度量，向量维度为1024
                .build();
        qdrantClient.createCollectionAsync("test-qdrant", vectorParams);//异步创建集合
    }

    /*
     往向量数据库新增文本记录
     */
    @GetMapping(value = "/embedding/add")//我们需要将文本数据向量化并存储到向量数据库中，以便后续进行相似性搜索
    public String add(){//作用：将一段文本转换为向量，并将其存储到Qdrant中
    //原理：首先将文本转换为TextSegment对象，并添加一些元数据（例如作者），然后使用嵌入模型将文本转换为向量，最后使用embeddingStore.add()方法将向量和文本片段存储起来。
        String prompt = """
                咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        TextSegment segment1 = TextSegment.from(prompt);
        segment1.metadata().put("author", "wgf1");//创建文本片段并添加元数据（作者信息）
        Embedding embedding1 = embeddingModel.embed(segment1).content();//embeddingModel.embed(prompt)：使用嵌入模型将文本转换为向量
        String result = embeddingStore.add(embedding1, segment1);//将向量和原始文本一起存储到向量数据库

        System.out.println(result);

        return result;
    }

    @GetMapping(value = "/embedding/query1")//为什么这么做：演示如何根据查询文本在向量数据库中找到相似的文本。
    public void query1(){// 作用：根据查询文本进行相似性搜索，返回最相似的文本片段。
        //原理：先将查询文本转换为向量，然后使用embeddingStore.search()方法进行搜索，设置最大返回结果数量。
        Embedding queryEmbedding = embeddingModel.embed("咏鸡说的是什么").content();//将查询文本"咏鸡说的是什么"转换为向量
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()//构建搜索请求
                .queryEmbedding(queryEmbedding)//设置查询向量
                .maxResults(1)//构建搜索请求，设置最大返回结果数为1
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);//执行搜索,在向量数据库中查找与查询向量最相似的文本
        System.out.println(searchResult.matches().get(0).embedded().text());//获取最相似结果的文本内容
    }

    @GetMapping(value = "/embedding/query2")//为什么这么做：演示如何在搜索时结合元数据过滤，以找到满足特定条件的相似文本。
    public void query2(){//作用：根据查询文本和元数据过滤器进行相似性搜索。
        //原理：与query1类似，但在搜索时添加了一个过滤器，要求元数据中的author字段等于"wgf2"。
        Embedding queryEmbedding = embeddingModel.embed("咏鸡").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("author").isEqualTo("wgf2"))//添加元数据（作者信息）
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
    }
}
